package com.example.batch.core.monitor;

import org.springframework.stereotype.Component;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Timer;
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.Gauge;
import com.zaxxer.hikari.HikariDataSource;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.Map;

@Component
@Slf4j
public class MetricsCollector {

    private final MeterRegistry registry;
    private final Map<String, Timer> queryTimers = new ConcurrentHashMap<>();
    private final Map<String, Counter> errorCounters = new ConcurrentHashMap<>();
    private final Map<String, Gauge> connectionGauges = new ConcurrentHashMap<>();
    
    public MetricsCollector(MeterRegistry registry) {
        this.registry = registry;
    }
    
    public void recordQueryTime(String dataSource, long timeMs) {
        getOrCreateTimer(dataSource).record(timeMs, TimeUnit.MILLISECONDS);
    }
    
    public void recordError(String dataSource, String errorType) {
        getOrCreateErrorCounter(dataSource, errorType).increment();
    }
    
    public void updateConnectionStats(String dataSource, HikariDataSource hikariDataSource) {
        Gauge.builder("datasource.connections.active", hikariDataSource,
                this::getActiveConnections)
            .tag("name", dataSource)
            .register(registry);
            
        Gauge.builder("datasource.connections.idle", hikariDataSource,
                this::getIdleConnections)
            .tag("name", dataSource)
            .register(registry);
            
        Gauge.builder("datasource.connections.pending", hikariDataSource,
                this::getPendingConnections)
            .tag("name", dataSource)
            .register(registry);
            
        // 添加更多连接池指标
        Gauge.builder("datasource.connections.timeout", hikariDataSource,
                this::getConnectionTimeouts)
            .tag("name", dataSource)
            .register(registry);
    }
    
    private Timer getOrCreateTimer(String dataSource) {
        return queryTimers.computeIfAbsent(dataSource,
            name -> Timer.builder("datasource.query.time")
                .tag("name", name)
                .register(registry));
    }
} 